// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/strings/string_view.h>
#include <pollux/common/base/simd_util.h>

namespace kumo::pollux {
    namespace {
        int32_t linearSearchSimple(
            StringView key,
            const StringView *strings,
            const int32_t *indices,
            int32_t numStrings) {
            if (indices) {
                for (auto i = 0; i < numStrings; ++i) {
                    if (strings[indices[i]] == key) {
                        return i;
                    }
                }
            } else {
                for (auto i = 0; i < numStrings; ++i) {
                    if (strings[i] == key) {
                        return i;
                    }
                }
            }
            return -1;
        }
    } // namespace

    // static
    int32_t StringView::linearSearch(
        StringView key,
        const StringView *strings,
        const int32_t *indices,
        int32_t numStrings) {
#if KSIMD_WITH_AVX2
        constexpr int64_t kBatch = ksimd::batch<uint64_t>::size;
        bool isInline = key.isInline();
        bool headOnly = key.size() <= 4;
        const char *body = key.data() + 4;
        int32_t bodySize = key.size() - 4;
        int32_t limit = numStrings & ~(kBatch - 1); // round down to full batches.
        if (indices) {
            uint64_t head = *reinterpret_cast<const uint64_t *>(&key);
            uint64_t inlined = reinterpret_cast<const uint64_t *>(&key)[1];
            ksimd::batch<int32_t, ksimd::sse2> indexVector;

            for (auto i = 0; i < limit; i += kBatch) {
                indexVector = simd::loadGatherIndices<uint64_t, int32_t>(indices + i)
                              << 1;
                auto heads =
                        simd::gather(reinterpret_cast<const uint64_t *>(strings), indexVector);
                uint16_t hits = simd::toBitMask(heads == head);
                if (LIKELY(!hits)) {
                    continue;
                }
                if (headOnly) {
                    return i + __builtin_ctz(hits);
                }
                while (hits) {
                    auto offset = bits::getAndClearLastSetBit(hits);
                    if (isInline
                            ? inlined ==
                              reinterpret_cast<const uint64_t *>(
                                  &strings[indices[i + offset]])[1]
                            : simd::memEqualUnsafe(
                                body,
                                strings[indices[i + offset]].data() + 4,
                                bodySize)) {
                        return i + offset;
                    }
                }
            }
            return linearSearchSimple(
                key, strings, indices + limit, numStrings - limit);
        } else {
            StringView key2[2];
            memcpy(&key2[0], &key, sizeof(key));
            memcpy(&key2[1], &key, sizeof(key));
            auto keyVector = ksimd::load_unaligned(reinterpret_cast<uint64_t *>(&key2));
            for (auto i = 0; i < limit; i += kBatch, strings += kBatch) {
                // Compare  4 StringViews  in 2 loads of 2 each.
                int32_t bits =
                        simd::toBitMask(
                            ksimd::load_unaligned(
                                reinterpret_cast<const uint64_t *>(strings)) == keyVector) |
                        (simd::toBitMask(
                             ksimd::load_unaligned(
                                 reinterpret_cast<const uint64_t *>(strings + 2)) == keyVector)
                         << 4);

                if ((bits & (1 + 4 + 16 + 64)) == 0) {
                    // Neither lane 0 or 2 or 4 or 6 hits
                    continue;
                }
                int offset = i;
                while (bits) {
                    auto low = bits & 3;
                    // At least first lane must match.
                    if (low & 1) {
                        // Both first and second lane match or only first word counts.
                        if (low == 3 || headOnly) {
                            return offset;
                        }
                        if (!isInline) {
                            if (simd::memEqualUnsafe(
                                body, strings[offset].data() + 4, bodySize)) {
                                return offset;
                            }
                        }
                    }
                    bits = bits >> 2;
                    ++offset;
                }
            }
            return linearSearchSimple(key, strings, nullptr, numStrings - limit);
        }
#else
  return linearSearchSimple(key, strings, indices, numStrings);
#endif
    }
} // namespace kumo::pollux
